package com.ruoyi.bearing.service.Impl;

import com.alibaba.fastjson2.JSONObject;
import com.ruoyi.bearing.domain.*;
import com.ruoyi.bearing.mapper.ModelTrainMapper;
import com.ruoyi.bearing.service.IModelInfoMangmentService;
import com.ruoyi.bearing.service.IModelTrainService;
import com.ruoyi.common.constant.FlaskModelConstants;
import com.ruoyi.common.enums.DiagnosisModelType;
import com.ruoyi.common.exception.flask.FlaskException;
import com.ruoyi.common.utils.file.FileUploadUtils;
import com.ruoyi.common.utils.flask.HttpFlask;
import com.ruoyi.common.utils.uuid.IdUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;

import java.io.IOException;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * 模型训练
 */
@Service
public class ModelTrainServiceImpl implements IModelTrainService {
    @Autowired
    IModelInfoMangmentService iModelInfoMangmentService;
    @Autowired
    ModelTrainMapper modelTrainMapper;

    /**
     * 诊断模型训练
     */
    @Override
    public String startDiagnosisModelTrain(MultipartFile file, ModelTrain modelTrain) throws IOException, IllegalAccessException {
        String res = "";
        modelTrain.setMid("diag" + IdUtils.simpleUUID());
        if (modelTrain.getModel().equals(DiagnosisModelType.SVM.getCode())) {
            modelTrain.setType(DiagnosisModelType.SVM.getInfo());
            res = modelToFlask(file, modelTrain, DiagnosisModelType.SVM.getCode());
        }
        if (modelTrain.getModel().equals(DiagnosisModelType.LDA.getCode())) {
            modelTrain.setType(DiagnosisModelType.LDA.getInfo());
            res = modelToFlask(file, modelTrain, DiagnosisModelType.LDA.getCode());
        }
        if (modelTrain.getModel().equals(DiagnosisModelType.CNN.getCode())) {
            modelTrain.setType(DiagnosisModelType.CNN.getInfo());
            res = modelToFlask(file, modelTrain, DiagnosisModelType.CNN.getCode());
        }
        if (modelTrain.getModel().equals(DiagnosisModelType.DBN.getCode())) {
            modelTrain.setType(DiagnosisModelType.DBN.getInfo());
            res = modelToFlask(file, modelTrain, DiagnosisModelType.DBN.getCode());
        }
        return res;
    }

    /**
     * 预测模型训练
     */
    /**
     * 预测模型训练
     */
    public String startPredictionModelTrain(MultipartFile file, ModelTrain modelTrain){

        return null;
    }

    /**
     * 将训练参数发送到Flask
     */
    private String modelToFlask(MultipartFile file, ModelTrain modelTrain, String type) throws IOException, IllegalAccessException {
        Object model = null;
        String csvFileName = "";
        String res = HttpFlask.sendGetFlask(FlaskModelConstants.Model_CSV_URL);
        if (!"".equals(res)) {
            JSONObject jsonObject = JSONObject.parseObject(res);
            String dir = (String) jsonObject.get("data");
            String csvName = FileUploadUtils.upload(dir, file);
            csvFileName = csvName.split("csv/")[1];
            if (type.equals(DiagnosisModelType.SVM.getCode())) {
                model = getModelSVM(modelTrain);
                ((ModelSVM) model).setFileName(csvFileName);
            }
            if (type.equals(DiagnosisModelType.LDA.getCode())) {
                model = getModelLDA(modelTrain);
                ((ModelLDA) model).setFileName(csvFileName);
            }
            if (type.equals(DiagnosisModelType.CNN.getCode())) {
                model = getModelCNN(modelTrain);
                ((ModelCNN) model).setFileName(csvFileName);
            }
            if (type.equals(DiagnosisModelType.DBN.getCode())) {
                model = getModelDBN(modelTrain);
                ((ModelDBN) model).setFileName(csvFileName);
            }
        } else {
            return "";
        }
        String postParam = HttpFlask.getKeyAndValueString(model);
        String resPost = HttpFlask.sendPostFlask(FlaskModelConstants.Model_DIAGNOSIS_URL, postParam);
        if (!"".equals(resPost) & !resPost.contains("请检查文件格式")) {
            int resInsert = setModelInfo(model, type);
            int resModelTrain = modelTrainMapper.insertModelTrain(modelTrain);
            if (0 == resInsert) {
                resPost = "";
            }
            if (0 == resModelTrain) {
                resPost = "";
            }
        }
        return resPost;
    }

    private ModelSVM getModelSVM(ModelTrain modelTrain) {
        ModelSVM model = new ModelSVM();
        model.setMid(modelTrain.getMid());
        model.setType(DiagnosisModelType.SVM.getInfo());
        model.setName(modelTrain.getName());
        model.setC(modelTrain.getC());
        model.setR(modelTrain.getR());
        return model;
    }

    private int setModelInfo(Object model, String type) {
        try {
            ModelInfoMangment modelInfoMangment = new ModelInfoMangment();
            if (type.equals(DiagnosisModelType.SVM.getCode())) {
                modelInfoMangment.setMid(((ModelSVM) model).getMid());
                modelInfoMangment.setName(((ModelSVM) model).getName());
                modelInfoMangment.setType(DiagnosisModelType.SVM.getInfo());
                modelInfoMangment.setState("0");
            }
            if (type.equals(DiagnosisModelType.LDA.getCode())) {
                modelInfoMangment.setMid(((ModelLDA) model).getMid());
                modelInfoMangment.setName(((ModelLDA) model).getName());
                modelInfoMangment.setType(DiagnosisModelType.LDA.getInfo());
                modelInfoMangment.setState("0");
            }
            if (type.equals(DiagnosisModelType.CNN.getCode())) {
                modelInfoMangment.setMid(((ModelCNN) model).getMid());
                modelInfoMangment.setName(((ModelCNN) model).getName());
                modelInfoMangment.setType(DiagnosisModelType.CNN.getInfo());
                modelInfoMangment.setState("0");
            }
            if (type.equals(DiagnosisModelType.DBN.getCode())) {
                modelInfoMangment.setMid(((ModelDBN) model).getMid());
                modelInfoMangment.setName(((ModelDBN) model).getName());
                modelInfoMangment.setType(DiagnosisModelType.DBN.getInfo());
                modelInfoMangment.setState("0");
            }
            return iModelInfoMangmentService.insertModelInfoMangment(modelInfoMangment);
        } catch (Exception e) {
            e.printStackTrace();
            return 0;
        }

    }

    private ModelLDA getModelLDA(ModelTrain modelTrain) {
        ModelLDA model = new ModelLDA();
        model.setMid(modelTrain.getMid());
        model.setType(DiagnosisModelType.LDA.getInfo());
        model.setName(modelTrain.getName());
        return model;
    }

    private ModelCNN getModelCNN(ModelTrain modelTrain) {
        ModelCNN model = new ModelCNN();
        model.setMid(modelTrain.getMid());
        model.setType(DiagnosisModelType.CNN.getInfo());
        model.setName(modelTrain.getName());
        model.setLr(modelTrain.getLr());
        model.setEpoch(modelTrain.getEpoch());
        model.setBatch(modelTrain.getBatch());
        model.setOptimizer(modelTrain.getOptimizer());
        model.setLayer(modelTrain.getLayer());
        model.setKernelSize(modelTrain.getKernelSize());
        return model;
    }

    private ModelDBN getModelDBN(ModelTrain modelTrain) {
        ModelDBN model = new ModelDBN();
        model.setMid(modelTrain.getMid());
        model.setType(DiagnosisModelType.DBN.getInfo());
        model.setName(modelTrain.getName());
        model.setLr(modelTrain.getLr());
        model.setEpoch(modelTrain.getEpoch());
        model.setBatch(modelTrain.getBatch());
        model.setOptimizer(modelTrain.getOptimizer());
        model.setLayer(modelTrain.getLayer());
        return model;
    }
}
